package step

import (
	"bytes"
	"errors"
	"github.com/wanchain/go-mpc/common/hexutil"
	"github.com/wanchain/go-mpc/log"
	"github.com/wanchain/go-mpc/p2p/discover"
	"github.com/wanchain/go-mpc/storeman/osmconf"
	"github.com/wanchain/go-mpc/storeman/schnorrcomm"
	schcomm "github.com/wanchain/go-mpc/storeman/schnorrcomm"
	mpcprotocol "github.com/wanchain/go-mpc/storeman/storemanmpc/protocol"
	"github.com/wanchain/go-mpc/storeman/validator"
	"math/big"
	"strconv"
	"sync"
	"time"
)

type MpcContextFunc interface {
	UptAllStepsWtCount(waitingCount int)
	InsertDropPeers([]discover.NodeID)
	GetInput() *mpcprotocol.SignInput
	GetAlg() uint8
}

type BaseGenerator struct {
	mpc MpcContextFunc
}

type BaseStep struct {
	msgChan       chan *mpcprotocol.StepMessage
	finish        chan error
	waiting       int
	waitAll       bool // true: wait all
	stepId        int
	notRecvPeers  map[discover.NodeID]*discover.NodeID
	needDropPeers map[discover.NodeID]bool
	mpcResult     mpcprotocol.MpcResultInterface
	schnorrMpcer  mpcprotocol.SchnorrMPCer
	curveType     uint16
	ContextId     uint64
	MpcContext    MpcContextFunc
	totalNum      int
	crtHndMutex   sync.Mutex // createMessage and HandleMessage mutext

	grpIdStr   string
	grpIdBytes []byte
}

func (step *BaseStep) GetMpcResult() mpcprotocol.MpcResultInterface {
	return step.mpcResult
}

func (step *BaseStep) InitStep(mpcResult mpcprotocol.MpcResultInterface) error {
	step.mpcResult = mpcResult
	return nil
}

func CreateBaseStep(totalNum int, wait int) *BaseStep {
	step := &BaseStep{
		msgChan: make(chan *mpcprotocol.StepMessage, totalNum+3),
		finish:  make(chan error, 3)}

	if wait >= 0 {
		step.waiting = wait
	} else {
		step.waiting = totalNum
	}
	step.waitAll = true
	step.totalNum = totalNum
	step.needDropPeers = make(map[discover.NodeID]bool)

	return step
}

func (step *BaseStep) InsertDropPeer(peer *discover.NodeID) error {
	step.needDropPeers[*peer] = true
	return nil
}

func (step *BaseStep) GetDropList() string {
	nodeIds := make([]*discover.NodeID, 0)
	for key, _ := range step.needDropPeers {
		nodeIds = append(nodeIds, &key)
	}
	return discover.ArrNodeIdpsToStr(nodeIds)
}

func (step *BaseStep) IsNeedDrop(peer *discover.NodeID) bool {
	ok, _ := step.needDropPeers[*peer]
	if ok {
		return true
	} else {
		return false
	}
}

func (step *BaseStep) InitMessageLoop(msger mpcprotocol.GetMessageInterface) error {
	log.SyslogInfo("BaseStep.InitMessageLoop begin")
	if step.waiting <= 0 {
		step.finish <- nil
	} else {
		go func() {
			log.SyslogInfo("InitMessageLoop begin")

			for {
				err := step.HandleMessage(msger)
				if err != nil {
					if err != mpcprotocol.ErrQuit {
						log.SyslogErr("BaseStep::InitMessageLoop", "InitMessageLoop fail, get message err, err", err.Error())
					}

					break
				}
			}
		}()
	}

	return nil
}

func (step *BaseStep) Quit(err error) {
	step.msgChan <- nil
	step.finish <- err
}

func (step *BaseStep) FinishStep() error {
	var timeOut time.Duration
	// firstStep
	if step.stepId == 0 {
		timeOut = mpcprotocol.MPCTimeOutValidate
	} else {
		timeOut = mpcprotocol.MPCTimeOut
	}
	select {
	case err := <-step.finish:
		if err != nil {
			log.SyslogErr("BaseStep::FinishStep", " get a step finish error. err", err.Error())
		}
		step.msgChan <- nil

		return err
	case <-time.After(timeOut):
		log.SyslogWarning("BaseStep.FinishStep, wait step finish timeout")

		step.ClearChannel(step.msgChan) // delete message which is time out
		step.msgChan <- nil

		if !step.waitAll {
			return nil
		}
		return mpcprotocol.ErrTimeOut
	}
}

func (step *BaseStep) ClearChannel(msgChan chan *mpcprotocol.StepMessage) {
	timeOut := time.Millisecond
	for {
		select {
		case v, ok := <-msgChan:
			if !ok {
				log.Info("channel closed")
				return
			}
			log.Info("drop message of channel", "msg", v.String())
		case <-time.After(timeOut):
			log.Info("no message to be dropped of channel")
			return
		}
	}
}

func (step *BaseStep) GetMessageChan() chan *mpcprotocol.StepMessage {
	return step.msgChan
}

func (step *BaseStep) HandleMessage(msger mpcprotocol.GetMessageInterface) error {
	var msg *mpcprotocol.StepMessage
	select {
	case msg = <-step.msgChan:
		if msg == nil {
			log.SyslogInfo("BaseStep get a quit msg")
			return mpcprotocol.ErrQuit
		}

		if msg.StepId != step.GetStepId() {
			log.SyslogErr("Get message is not in the right steps",
				"should step", step.stepId,
				"receive step", msg.StepId)
		} else {
			if len(step.needDropPeers) != 0 {
				log.SyslogInfo("BaseStep::HandleMessage", "step message", msg.String(), "drop list", step.GetDropList(), "ctx", step.GetMpcContextId())
			}
			if step.waiting > 0 && !step.IsNeedDrop(msg.PeerID) && !step.CheckMaliceSm(msg) && step.checkSndIndex(msg) {

				defer step.crtHndMutex.Unlock()
				step.crtHndMutex.Lock()

				if msger.HandleMessage(msg) {
					delete(step.notRecvPeers, *msg.PeerID)

					step.waiting--
					if step.waiting <= 0 {
						step.finish <- nil
					}
				}

			}
		}
	}

	return nil
}

func (step *BaseStep) CheckMaliceSm(msg *mpcprotocol.StepMessage) bool {
	if !schnorrcomm.OpenSlsh {
		return false
	}
	_, grpIdString, _ := osmconf.GetGrpId(step.mpcResult)
	inx, _ := osmconf.GetOsmConf().GetInxByNodeId(grpIdString, msg.PeerID)
	b, err := validator.IsMalice(grpIdString, inx)
	if err != nil {
		log.SyslogWarning("BaseStep::CheckMalicSm", "grpId", grpIdString, "index", inx, "error", err.Error())
	}
	log.SyslogWarning("CheckMaliceSm", "result", b, "grpId", grpIdString, "index", inx)
	return b
}

func (step *BaseStep) SetWaitAll(waitAll bool) {
	step.waitAll = waitAll
}

func (step *BaseStep) SetWaiting(waiting int) {
	log.Info("BaseStep", "SetWaiting waiting count", waiting, "original waiting count", step.waiting, "contexId", step.ContextId, "stepId", step.stepId)
	if step.waiting > waiting {
		step.waiting = waiting
	}
}

func (step *BaseStep) SetStepId(stepId int) {
	step.stepId = stepId
}

func (step *BaseStep) GetStepId() uint8 {
	return uint8(step.stepId)
}

func (step *BaseStep) SetSchnorrMpcer(smcer mpcprotocol.SchnorrMPCer) {
	step.schnorrMpcer = smcer
}

func (step *BaseStep) SchnorrMpcer() mpcprotocol.SchnorrMPCer {
	return step.schnorrMpcer
}

func (step *BaseStep) GetMsgGens() []MpcMessageGenerator {
	return nil
}

func (step *BaseStep) SetCurveType(ct uint16) {
	step.curveType = ct
}

func (step *BaseStep) CurveType() uint16 {
	return step.curveType
}

func (step *BaseStep) TotalNum() int {
	return step.totalNum
}

func (step *BaseStep) SetMpcContextId(contextId uint64) {
	step.ContextId = contextId
}

func (step *BaseStep) GetMpcContextId() uint64 {
	return step.ContextId
}

func (step *BaseStep) GetNodeIdsNotInColletion(grpIdString string, rcvColInter *big.Int) ([]discover.NodeID, int, error) {
	ret := make([]discover.NodeID, 0)
	count := 0
	totalNumber, err := osmconf.GetOsmConf().GetTotalNum(grpIdString)
	if err != nil {
		return nil, 0, err
	}
	for i := 0; i < int(totalNumber); i++ {
		b, err := osmconf.IsHaming(rcvColInter, uint16(i))
		if err != nil {
			return nil, 0, err
		}
		if !b {
			nodeId, err := osmconf.GetOsmConf().GetNodeIdByIndex(grpIdString, uint16(i))
			if err != nil {
				return nil, 0, err
			}
			ret = append(ret, *nodeId)
			count++
		}
	}
	return ret, int(totalNumber) - count, nil
}

func (step *BaseStep) InsertDropPeers(nodes []discover.NodeID) {
	for _, id := range nodes {
		step.InsertDropPeer(&id)
	}
}

func (step *BaseStep) SetMpcContext(mpc MpcContextFunc) {
	step.MpcContext = mpc
}

func (step *BaseStep) GetMpcContext() MpcContextFunc {
	return step.MpcContext
}

func (msStep *BaseStep) checkContent(sshare, m *big.Int, rpkShare, gpkShare mpcprotocol.CurvePointer) (bool, error) {

	if sshare == nil || m == nil {
		return false, errors.New("sshare is nil or m is nil")
	}
	smpcer := msStep.schnorrMpcer

	if !smpcer.IsOnCurve(rpkShare) || !smpcer.IsOnCurve(gpkShare) {
		return false, errors.New("rpkShare or gpkShare is invalid pk or gpkShare is invalid pk")
	}

	algBytes := msStep.GetMpcContext().GetInput().Alg
	alg := int(big.NewInt(0).SetBytes(algBytes).Uint64())
	if alg == mpcprotocol.ALGSCHNORR {
		sshareG, _ := smpcer.SkG(sshare)
		mPkShare, _ := smpcer.MulPK(m, gpkShare)

		if !smpcer.IsOnCurve(mPkShare) {
			return false, errors.New("mPkShare is invalid pk or gpkShare is invalid pk")
		}

		pkTemp, err := smpcer.Add(rpkShare, mPkShare)
		if err != nil {
			return false, errors.New("add rpkShare mPkShare error")
		}
		left := sshareG
		right := pkTemp
		return smpcer.Equal(left, right), nil
	} else if alg == mpcprotocol.ALGSCHNORR340 {

		xBytesGpkShare, err := smpcer.XBytes(gpkShare)
		if err != nil {

			log.SyslogErr("MpcSSahreJudgeStepUni::checkContent", "xBytesGpkShare err", err.Error(), "ctx", msStep.GetMpcContextId())
			return false, err

		}

		xBytesRpkShare, err := smpcer.XBytes(rpkShare)
		if err != nil {

			log.SyslogErr("MpcSSahreJudgeStepUni::checkContent", "xBytesRpkShare err", err.Error(), "ctx", msStep.GetMpcContextId())
			return false, err

		}

		nSubE, err := smpcer.NSubX(m)
		if err != nil {
			log.SyslogErr("MpcSSahreJudgeStepUni::checkContent", "NSubX err", err.Error(), "ctx", msStep.GetMpcContextId())
			return false, err
		}

		// ssG + (n-e)*gpkLift = rpk
		ssG, err := smpcer.SkG(sshare)
		if err != nil {
			log.SyslogErr("MpcSSahreJudgeStepUni::checkContent", "SkG err", err.Error(), "ctx", msStep.GetMpcContextId())
			return false, err
		}

		gpkLift, err := smpcer.LiftX(big.NewInt(0).SetBytes(xBytesGpkShare))
		if err != nil {
			log.SyslogErr("MpcSSahreJudgeStepUni::checkContent", "LiftX err", err.Error(), "ctx", msStep.GetMpcContextId())
			return false, err
		}

		//gpkLift*(n-e)
		temp, err := smpcer.MulPK(nSubE, gpkLift)
		if err != nil {
			log.SyslogErr("MpcSSahreJudgeStepUni::checkContent", "MulPK err", err.Error(), "ctx", msStep.GetMpcContextId())
			return false, err
		}

		temp1, err := smpcer.MulPK(nSubE, gpkShare)
		if err != nil {
			log.SyslogErr("MpcSSahreJudgeStepUni::checkContent", "MulPK err", err.Error(), "ctx", msStep.GetMpcContextId())
			return false, err
		}

		// left
		left, err := smpcer.Add(ssG, temp)
		if err != nil {
			log.SyslogErr("MpcSSahreJudgeStepUni::checkContent", "Add err", err.Error(), "ctx", msStep.GetMpcContextId())
			return false, err
		}

		left1, err := smpcer.Add(ssG, temp1)
		if err != nil {
			log.SyslogErr("MpcSSahreJudgeStepUni::checkContent", "Add err", err.Error(), "ctx", msStep.GetMpcContextId())
			return false, err
		}

		if !smpcer.IsOnCurve(left) || !smpcer.IsOnCurve(left1) {
			log.SyslogErr("MpcSSahreJudgeStepUni::checkContent left is Not on curve", "ctx", msStep.GetMpcContextId())
			return false, mpcprotocol.ErrNotOnCurve
		}

		xBytesLeft, err := smpcer.XBytes(left)
		if err != nil {
			log.SyslogErr("MpcSSahreJudgeStepUni::checkContent", "XBytes xBytesLeft err", err.Error(), "ctx", msStep.GetMpcContextId())
			return false, err
		}

		xBytesLeft1, err := smpcer.XBytes(left1)
		if err != nil {
			log.SyslogErr("MpcSSahreJudgeStepUni::checkContent", "XBytes xBytesLeft err", err.Error(), "ctx", msStep.GetMpcContextId())
			return false, err
		}

		if bytes.Equal(xBytesLeft, xBytesRpkShare) || bytes.Equal(xBytesLeft1, xBytesRpkShare) {
			return true, nil
		} else {
			return false, mpcprotocol.ErrVerifyFailed
		}
	}
	return true, nil
}

func (msStep *BaseStep) getRPkShare(index uint16, sigIndex int) (mpcprotocol.CurvePointer, error) {

	key := mpcprotocol.RPkShare + strconv.Itoa(int(index)) + strconv.Itoa(sigIndex)

	log.SyslogInfo("getRPkShare", "index", index, "key", key, "ctx", msStep.GetMpcContextId())
	pkBytes, err := msStep.mpcResult.GetByteValue(key)
	if err != nil {
		log.SyslogErr("getRPkShare", "err", err.Error(), "ctx", msStep.GetMpcContextId())
	}
	return msStep.schnorrMpcer.UnMarshPt(pkBytes)
}

func (msStep *BaseStep) getm(sigIndex int) (*big.Int, error) {
	// check signVerify
	result := msStep.mpcResult
	hashMBytes := msStep.GetMpcContext().GetInput().HashData[sigIndex]
	// rpk : R
	key := mpcprotocol.RPk + strconv.Itoa(sigIndex)
	rpkBytes, _ := result.GetByteValue(key)
	log.SyslogDebug("MpcSStep getm", "rpkBytes", hexutil.Encode(rpkBytes), "ctx", msStep.GetMpcContextId())

	var m *big.Int
	smpcer := msStep.SchnorrMpcer()
	algBytes := msStep.GetMpcContext().GetInput().Alg
	alg := int(big.NewInt(0).SetBytes(algBytes).Uint64())

	curveIdBytes := msStep.GetMpcContext().GetInput().Curve
	curveId := int(big.NewInt(0).SetBytes(curveIdBytes).Uint64())

	if alg == mpcprotocol.ALGSCHNORR {
		if curveId == mpcprotocol.SK256Curve {
			m, _ = schcomm.GetSchnorrSecE(hashMBytes, rpkBytes, msStep.GetMpcContext().GetInput().PKBytes, smpcer.GetMod())
		} else {
			m, _ = schcomm.GetSchnorrE(hashMBytes, rpkBytes, smpcer.GetMod())
		}

	} else if alg == mpcprotocol.ALGSCHNORR340 {
		xbytesRGPK, err := smpcer.XBytesByBytes(rpkBytes)
		if err != nil {
			return nil, err
		}
		xbytesGPK, err := smpcer.XBytesByBytes(msStep.MpcContext.GetInput().PKBytes)
		if err != nil {
			return nil, err
		}
		m, _ = schcomm.GetSchnorr340E(hashMBytes, xbytesRGPK, xbytesGPK, smpcer.GetMod())
	}
	return m, nil
}

func (msStep *BaseStep) getGPKShare(index uint16) (mpcprotocol.CurvePointer, error) {

	_, grpIdString, err := osmconf.GetGrpId(msStep.mpcResult)
	if err != nil {
		return nil, err
	}

	// todo should get from input??
	_, gpkStr, err := osmconf.GetGrpGpk(msStep.mpcResult)
	if err != nil {
		return nil, err
	}
	log.SyslogInfo("MpcSStep getGPKShare", "grpIdString", grpIdString, "index", index, "CurveType", msStep.CurveType(), "ctx", msStep.GetMpcContextId())
	gpkShareBytes, err := osmconf.GetOsmConf().GetPKShareBytes(grpIdString, gpkStr, index, msStep.CurveType())
	if err != nil {
		return nil, err
	}
	return msStep.schnorrMpcer.UnMarshPt(gpkShareBytes)
}

func (msStep *BaseStep) checkSndIndex(msg *mpcprotocol.StepMessage) bool {
	_, grpIdString, _ := msStep.GetGrpId()
	_, err := osmconf.GetOsmConf().GetInxByNodeId(grpIdString, msg.PeerID)
	if err != nil {
		log.SyslogInfo("BaseStep::checkSndIndex", "err", err.Error(), "ctx", msStep.GetMpcContextId())
		return false
	}
	return true
}

func (msStep *BaseStep) GetGrpId() ([]byte, string, error) {
	if len(msStep.grpIdBytes) != 0 && len(msStep.grpIdStr) != 0 {
		return msStep.grpIdBytes, msStep.grpIdStr, nil
	} else {
		grpIdBytes, grpIdStr, err := osmconf.GetGrpId(msStep.GetMpcResult())
		msStep.grpIdBytes = grpIdBytes
		msStep.grpIdStr = grpIdStr
		return msStep.grpIdBytes, msStep.grpIdStr, err
	}
}
